#include "stdafx.h"
#include "MoveGenerator.h"
#include "MoveLookups.h"

using namespace std;




MoveGenerator::MoveGenerator()
{
	
}

void MoveGenerator::Generate(Board &b, Move *lastMove, std::vector<Move> &moves)
{
	Generate(b, lastMove, moves, true);
}

void MoveGenerator::Generate(Board &b, Move *lastMove, std::vector<Move> &moves, bool checkForMate)
{
	ClearAttackBoard();

	for (byte i = 0; i < 64; ++i)
	{
		Piece p = b.Squares[i];

		if (!p)
			continue;

		switch (p.Type)
		{
		case PieceType::Pawn:
			GeneratePawnMoves(b, i, moves);
			break;
		case PieceType::Rook:
			GenerateRookMoves(b, i, moves);
			break;
		case PieceType::Knight:
			GenerateKnightMoves(b, i, moves);
			break;
		case PieceType::Bishop:
			GenerateBishopMoves(b, i, moves);
			break;
		case PieceType::Queen:
			GenerateQueenMoves(b, i, moves);
			break;
		case PieceType::King:
			// Since the available moves for the king depends on 
			// which squares are safe, defer move generation until the end
			if (p.Color == PieceColor::Black)
				_blackKingPosition = i;
			else
				_whiteKingPosition = i;
			break;
		}
	}

	byte kp1 = _whiteKingPosition, kp2 = _blackKingPosition;
	if (b.WhoseMove == PieceColor::White)
		swap(kp1, kp2);

	// Generate moves for the opposing king first so that we can mark attacked squares
	GenerateKingMoves(b, kp1, moves);
	GenerateKingMoves(b, kp2, moves);

	if (b.WhoseMove == PieceColor::Black)
		GenerateCastleMoves(b, PieceColor::Black, moves);
	else
		GenerateCastleMoves(b, PieceColor::White, moves);

	//if (checkForMate && (_blackCheck || _whiteCheck || moves.empty()))
	if (checkForMate)
	{
		SearchForMate(b, moves);

		/*if (lastMove)
		{
		if (b.WhoseMove == PieceColor::White)
		lastMove->NumWhiteMoves = moves.size();
		else
		lastMove->NumBlackMoves = moves.size();

		if (moves.empty())
		lastMove->Flags |= STALEMATE;
		}

		if (moves.empty())
		b.CurrentState |= STALEMATE;*/
	}
	else
	{
		UpdateFlags(b.CurrentState);

		UpdateFlags(lastMove->Flags);

		if (b.WhoseMove == PieceColor::White)
			lastMove->NumBlackMoves = moves.size();
		else
			lastMove->NumWhiteMoves = moves.size();
	}

	/*UpdateFlags(b.CurrentState);

	if (moves.empty() && !_blackMate && !_whiteMate)
	{
		b.CurrentState |= STALEMATE;
		if (lastMove)
			lastMove->Flags |= STALEMATE;
	}*/

	/*if (lastMove)
	{
		UpdateFlags(lastMove->Flags);
		if (b.WhoseMove == PieceColor::Black)
			lastMove->NumBlackMoves = moves.size();
		else
			lastMove->NumWhiteMoves = moves.size();
	}*/
}

void MoveGenerator::GeneratePawnMoves(Board &b, byte idx, std::vector<Move> &moves)
{
	// TODO: Support En Passant

	Piece p = b.Squares[idx];

	sbyte sidx = idx;
	sbyte row = (sidx & ~0x7) >> 3;
	sbyte homeRow = p.Color == PieceColor::Black ? 1 : 6;
	sbyte delta = p.Color == PieceColor::Black ? 8 : -8;

	sbyte col = sidx & 0x7;

	uint flag = (row == 1 && delta < 0) || (row == 6 && delta > 0) ? PROMOTED : NONE;

	sbyte firstPos = sidx + delta;

	if (firstPos < 0 || firstPos >= 64)
		return;

	bool canForward1 = true;
	if (!b.Squares[firstPos])
	{
		canForward1 = true;

		if (p.Color == b.WhoseMove && _generateMoves)
			moves.emplace_back(p, idx, firstPos, flag);
	}

	sbyte secondPos = firstPos + delta;

	if (canForward1 && secondPos >= 0 && secondPos < 64 && !b.Squares[secondPos] && row == homeRow)
	{
		if (p.Color == b.WhoseMove && _generateMoves)
			moves.emplace_back(p, idx, secondPos);
	}

	// Test to see if we can attack to the left
	if (col > 0 && b.Squares[firstPos - 1])
	{
		AnalyzeMove(b, idx, firstPos - 1, moves, flag);
	}

	// Test to see if we can attack to the right
	if (col < 7 && b.Squares[firstPos + 1])
	{
		AnalyzeMove(b, idx, firstPos + 1, moves, flag);
	}
}

void MoveGenerator::GenerateRookMoves(Board &b, byte idx, std::vector<Move> &moves)
{
	sbyte sIdx = (sbyte) idx;

	StateFlags moveFlag = NONE;
	switch (idx)
	{
	case 0:
		moveFlag = B_L_ROOK_MOVED;
		break;
	case 7:
		moveFlag = B_R_ROOK_MOVED;
		break;
	case 56:
		moveFlag = W_L_ROOK_MOVED;
		break;
	case 63:
		moveFlag = W_R_ROOK_MOVED;
		break;
	}

	auto gen = [&, this](byte dstPos) -> bool
	{
		return AnalyzeMove(b, idx, dstPos, moves, moveFlag);
	};

	sbyte sRow = idx & ~0x7;
	byte row = idx & ~0x7;

	for (sbyte curr = sIdx - 1; curr >= row && gen(curr); --curr);
	for (byte curr = idx + 1; curr < row + 8 && gen(curr); ++curr);
	for (sbyte curr = sIdx - 8; curr >= 0 && gen(curr); curr -= 8);
	for (byte curr = idx + 8; curr < NUM_SQUARES && gen(curr); curr += 8);
}

void MoveGenerator::GenerateKnightMoves(Board &b, byte idx, std::vector<Move> &moves)
{
	const byte *pLookup = &s_possibleKnightMoves.Moves[idx * KnightMoves::RowSize];

	for (; *pLookup != 64; ++pLookup)
	{
		AnalyzeMove(b, idx, *pLookup, moves);
	}
}

void MoveGenerator::GenerateBishopMoves(Board &b, byte idx, std::vector<Move> &moves)
{
	const byte *pLookup = s_possibleBishopMoves.Moves + idx * BishopMoves::RowSize;

	int left = 4;

	for (int left = 4; left > 0; --left)
	{
		byte ct = *pLookup++;

		for (byte i = 0; i < ct; ++i)
		{
			if (!AnalyzeMove(b, idx, pLookup[i], moves))
				break;
		}

		pLookup += ct;
	}
}

void MoveGenerator::GenerateQueenMoves(Board &b, byte idx, std::vector<Move> &moves)
{
	GenerateRookMoves(b, idx, moves);
	GenerateBishopMoves(b, idx, moves);
}

void MoveGenerator::GenerateKingMoves(Board &b, byte idx, std::vector<Move> &moves)
{
	Piece p = b.Squares[idx];

	bool *attackBoard = nullptr;
	switch (p.Color)
	{
	case PieceColor::White:
		attackBoard = _blackAttackBoard;
		break;
	case PieceColor::Black:
		attackBoard = _whiteAttackBoard;
		break;
	}

	StateFlags moveFlag = p.Color == PieceColor::Black ? B_KING_MOVED : W_KING_MOVED;

	byte *pLookup = s_possibleKingMoves.Moves + idx * KingMoves::RowSize;
	
	byte numPositions = *pLookup++;

	for (byte *pEnd = pLookup + numPositions; pLookup != pEnd; ++pLookup)
	{
		if (!attackBoard[*pLookup])
			AnalyzeMove(b, idx, *pLookup, moves, moveFlag);
	}
}

void MoveGenerator::GenerateCastleMoves(Board &b, PieceColor color, std::vector<Move> &moves)
{
	// The flags that cannot be set for a castle to be valid
	uint leftCastle = B_L_ROOK_MOVED | B_KING_MOVED | BLACK_CASTLED,
		 rightCastle = B_R_ROOK_MOVED | B_KING_MOVED | BLACK_CASTLED;
	byte kIdx = _blackKingPosition;

	bool *attacking = _whiteAttackBoard;

	if (color == PieceColor::White)
	{
		leftCastle = W_L_ROOK_MOVED | W_KING_MOVED | WHITE_CASTLED;
		rightCastle = W_R_ROOK_MOVED | W_KING_MOVED | WHITE_CASTLED;
		kIdx = _whiteKingPosition;
		attacking = _blackAttackBoard;
	}

	// Cannot castle if the king is under attack
	if (attacking[kIdx])
		return;

	Piece king = b.Squares[kIdx];

	if (!(b.CurrentState & leftCastle))
	{
		byte rIdx = kIdx - 4;

		if (!attacking[rIdx] && !attacking[kIdx - 2])
		{
			bool blank = true;
			for (byte i = rIdx + 1; i < kIdx; ++i)
			{
				if (b.Squares[i])
				{
					blank = false;
					break;
				}
			}

			if (blank)
			{
				moves.emplace_back(king, kIdx, kIdx - 2, rIdx, kIdx - 1, MoveType::Castle, leftCastle | CASTLED);
			}
		}
	}

	if (!(b.CurrentState & rightCastle))
	{
		byte rIdx = kIdx + 3;

		if (!attacking[rIdx] && !attacking[kIdx + 2])
		{
			bool blank = true;
			for (byte i = kIdx + 1; i < rIdx; ++i)
			{
				if (b.Squares[i])
				{
					blank = false;
					break;
				}
			}

			if (blank)
			{
				moves.emplace_back(king, kIdx, kIdx + 2, rIdx, kIdx + 1, MoveType::Castle, rightCastle | CASTLED);
			}
		}
	}
}

bool MoveGenerator::AnalyzeMove(Board &b, byte srcPos, byte dstPos, std::vector<Move> &moves, uint flags /*= NONE*/)
{
	Piece pcMoving = b.Squares[srcPos];

	bool *pCheckColor = nullptr;

	switch (pcMoving.Color)
	{
	case PieceColor::Black:
		_blackAttackBoard[dstPos] = true;
		pCheckColor = &_whiteCheck;
		break;
	case PieceColor::White:
		_whiteAttackBoard[dstPos] = true;
		pCheckColor = &_blackCheck;
		break;
	}

	Piece attackedPiece;

	// If the square is empty, then the move is always valid
	if (!(attackedPiece = b.Squares[dstPos]))
	{
		if (b.WhoseMove == pcMoving.Color && _generateMoves)
			moves.emplace_back(pcMoving, srcPos, dstPos, flags);
		return true;
	}

	// Obviously cannot capture own pieces
	if (attackedPiece.Color == pcMoving.Color)
		return false;

	switch (attackedPiece.Type)
	{
		// Mark the king to be in check
	case PieceType::King:
		*pCheckColor = true;
		break;
	}

	if (b.WhoseMove == pcMoving.Color && _generateMoves)
		moves.emplace_back(pcMoving, attackedPiece, srcPos, dstPos, flags);
	return false;
}

void MoveGenerator::UpdateFlags(uint &flags)
{
	if (_blackCheck)
		flags |= BLACK_CHECK;
	if (_whiteCheck)
		flags |= WHITE_CHECK;
	if (_blackMate)
		flags |= BLACK_MATE;
	if (_whiteMate)
		flags |= WHITE_MATE;
}

void MoveGenerator::SearchForMate(Board &b, std::vector<Move> &moves)
{
	PieceColor movingSide = b.WhoseMove;

	bool foundNonCheckBlack = false;
	bool foundNonCheckWhite = false;

	bool blackCheck = _blackCheck,
		whiteCheck = _whiteCheck;

	uint oldState = b.CurrentState;

	//_generateMoves = false;
	std::vector<Move> tmpMoves;

	for (int i = moves.size() - 1; i >= 0; --i)
	{
		Move &move = moves[i];

		// Reset the check flags every time
		_blackCheck = false;
		_whiteCheck = false;

		b.DoMove(move);

		std::vector<Move> nextMoves;
		Generate(b, &move, nextMoves, false);

		bool b1check = _blackCheck,
			 w1check = _whiteCheck;

		//bool blackCheck = b.IsBlackCheck(),
		//	whiteCheck = b.IsWhiteCheck();

		// If that move put the opposing player in check, then we need to analyze the other players responses
		// to see if it was a mate
		if ((_blackCheck && movingSide == PieceColor::White) ||
			(_whiteCheck && movingSide == PieceColor::Black))
		{
			uint mateFlag = (_blackCheck && movingSide == PieceColor::White) ? BLACK_MATE : WHITE_MATE;

			for (int nextI = nextMoves.size() - 1; nextI >= 0; --nextI)
			{
				Move &nextMove = nextMoves[nextI];

				_blackCheck = false;
				_whiteCheck = false;

				b.DoMove(nextMove);

				_generateMoves = false;
				Generate(b, &nextMove, tmpMoves, false);
				_generateMoves = true;

				bool eraseMove = (_blackCheck && movingSide == PieceColor::White) ||
								 (_whiteCheck && movingSide == PieceColor::Black);

				b.UndoMove(nextMove);

				if (eraseMove)
					nextMoves.erase(nextMoves.begin() + nextI);
			}

			if (nextMoves.empty())
			{
				move.Flags |= mateFlag;
			}
		}

		// This move is not valid if it would put the king in check
		bool eraseMove = (b1check && movingSide == PieceColor::Black) ||
						 (w1check && movingSide == PieceColor::White);

		b.UndoMove(move);

		if (eraseMove)
			moves.erase(moves.begin() + i);
	}
}

void MoveGenerator::ClearAttackBoard()
{
	memset(_blackAttackBoard, 0, sizeof(bool) * NUM_SQUARES);
	memset(_whiteAttackBoard, 0, sizeof(bool) * NUM_SQUARES);
}




